"""
Token estimation utilities for different AI models.
Provides more accurate token counting than simple word splitting.
"""

import re
from typing import Dict, Any


class TokenEstimator:
    """
    Token estimation for different AI models.
    Uses heuristics based on model-specific characteristics.
    """
    
    # 模型特定的token比率（基于经验值）
    MODEL_RATIOS = {
        "gpt-4": 0.75,           # GPT-4通常更高效
        "gpt-3.5-turbo": 0.8,    # GPT-3.5稍微冗余
        "claude": 0.75,          # Claude类似GPT-4
        "qwen": 1.2,             # 中文模型，字符更多
        "ernie": 1.2,            # 百度文心，中文优化
        "glm": 1.15,             # 智谱，中英混合
        "default": 1.0           # 默认值
    }
    
    @staticmethod
    def estimate_tokens(text: str, model_name: str = None) -> int:
        """
        估算文本的token数量。
        
        Args:
            text: 输入文本
            model_name: 模型名称（用于特定调整）
            
        Returns:
            估算的token数量
        """
        if not text:
            return 0
            
        # 基础计算：分词
        # 英文按空格和标点分词
        english_tokens = len(re.findall(r'\b\w+\b', text))
        
        # 中文按字符计算（每个汉字约1.5个token）
        chinese_chars = len(re.findall(r'[\u4e00-\u9fff]', text))
        chinese_tokens = chinese_chars * 1.5
        
        # 标点和特殊字符
        punctuation = len(re.findall(r'[^\w\s\u4e00-\u9fff]', text))
        
        # 基础token数
        base_tokens = english_tokens + chinese_tokens + punctuation * 0.5
        
        # 根据模型调整
        if model_name:
            # 查找最匹配的模型比率
            ratio = TokenEstimator.MODEL_RATIOS.get("default")
            for key in TokenEstimator.MODEL_RATIOS:
                if key in model_name.lower():
                    ratio = TokenEstimator.MODEL_RATIOS[key]
                    break
        else:
            ratio = TokenEstimator.MODEL_RATIOS["default"]
            
        # 应用模型特定比率
        estimated_tokens = int(base_tokens * ratio)
        
        # 最小值保护（至少按字符数的1/4）
        min_tokens = len(text) // 4
        
        return max(estimated_tokens, min_tokens)
    
    @staticmethod
    def estimate_cost(tokens: int, model_name: str, is_input: bool = True) -> float:
        """
        估算token成本（美元）。
        
        Args:
            tokens: token数量
            model_name: 模型名称
            is_input: 是否为输入token（否则为输出）
            
        Returns:
            估算的成本（美元）
        """
        # 模型成本（每1000 tokens，美元）
        # 注意：这些是示例价格，需要根据实际定价更新
        MODEL_COSTS = {
            "gpt-4": {"input": 0.03, "output": 0.06},
            "gpt-3.5-turbo": {"input": 0.001, "output": 0.002},
            "claude": {"input": 0.01, "output": 0.03},
            "qwen": {"input": 0.0008, "output": 0.002},
            "ernie": {"input": 0.001, "output": 0.002},
            "glm": {"input": 0.001, "output": 0.002},
        }
        
        # 查找模型成本
        costs = None
        for key in MODEL_COSTS:
            if key in model_name.lower():
                costs = MODEL_COSTS[key]
                break
                
        if not costs:
            # 默认成本
            costs = {"input": 0.002, "output": 0.004}
            
        # 计算成本
        cost_per_token = costs["input" if is_input else "output"] / 1000
        return tokens * cost_per_token